#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import numpy as np
import torch
from numpy import ndarray
from tqdm import tqdm
from torch import Tensor
from numpy import ndarray
from PIL.Image import Image as PILImage

from typing import NamedTuple, Tuple
from scene.cameras import Camera
from utils.graphics_utils import fov2focal
from utils.image_utils import soft_split_freqs_torch, soft_linear_split_freq, soft_split_numpy

WARNED = False

def loadCam(args, id, cam_info, resolution_scale):
    orig_w, orig_h = cam_info.image.size

    if args.resolution in [1, 2, 4, 8]:
        resolution = round(orig_w/(resolution_scale * args.resolution)), round(orig_h/(resolution_scale * args.resolution))
    else:  # should be a type that converts to float
        if args.resolution == -1:
            if orig_w > 1600:
                global WARNED
                if not WARNED:
                    print("[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.\n "
                        "If this is not desired, please explicitly specify '--resolution/-r' as 1")
                    WARNED = True
                global_down = orig_w / 1600
            else:
                global_down = 1
        else:
            global_down = orig_w / args.resolution

        scale = float(global_down) * float(resolution_scale)
        resolution = (int(orig_w / scale), int(orig_h / scale))

    resized_image_rgb = PILtoTorch(cam_info.image, resolution, args)
    gt_image = [rgb[:3, ...] for rgb in resized_image_rgb]    
    loaded_mask = None

    if resized_image_rgb[0].shape[1] == 4:
        loaded_mask = resized_image_rgb[0][3:4, ...]

    return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, 
                  FoVx=cam_info.FovX, FoVy=cam_info.FovY, 
                  image=gt_image, gt_alpha_mask=loaded_mask,
                  image_name=cam_info.image_name, uid=id, data_device=args.data_device)

def cameraList_from_camInfos(cam_infos, resolution_scale, args):
    camera_list = []

    for id, c in tqdm(enumerate(cam_infos)):
        camera_list.append(loadCam(args, id, c, resolution_scale))

    return camera_list

def camera_to_JSON(id, camera : Camera):
    Rt = np.zeros((4, 4))
    Rt[:3, :3] = camera.R.transpose()
    Rt[:3, 3] = camera.T
    Rt[3, 3] = 1.0

    W2C = np.linalg.inv(Rt)
    pos = W2C[:3, 3]
    rot = W2C[:3, :3]
    serializable_array_2d = [x.tolist() for x in rot]
    camera_entry = {
        'id' : id,
        'img_name' : camera.image_name,
        'width' : camera.width,
        'height' : camera.height,
        'position': pos.tolist(),
        'rotation': serializable_array_2d,
        'fy' : fov2focal(camera.FovY, camera.height),
        'fx' : fov2focal(camera.FovX, camera.width)
    }
    return camera_entry


class CameraInfo(NamedTuple):
    uid: int
    R: ndarray
    T: ndarray
    FovY: ndarray
    FovX: ndarray
    image: ndarray
    image_path: str
    image_name: str
    width: int
    height: int

    def to_json(self, id:int=None) -> dict:
        Rt = np.zeros((4, 4))
        Rt[:3, :3] = self.R.transpose()
        Rt[:3, 3] = self.T
        Rt[3, 3] = 1.0
        W2C = np.linalg.inv(Rt)
        pos = W2C[:3, 3]
        rot = W2C[:3, :3]
        return {
            'id': id,
            'img_name': self.image_name,
            'width': self.width,
            'height': self.height,
            'position': pos.tolist(),
            'rotation': [x.tolist() for x in rot],
            'fy': fov2focal(self.FovY, self.height),
            'fx': fov2focal(self.FovX, self.width)
        }

def PILtoTorch(img:PILImage, resize:Tuple[int, int]=None) -> Tensor:
    if resize: img = img.resize(resize)
    X: Tensor = torch.from_numpy(np.array(img)) / 255.0
    if img.mode == 'L': X = X.unsqueeze(dim=-1)
    return X.permute(2, 0, 1)

def load_camera(hp, training, resolution:float, id:int, cam_info:CameraInfo, resolution_scale:float) -> Camera:
    global WARNED
    
    orig_w, orig_h = cam_info.image.size
    if resolution in [1, 2, 4, 8]:
        resolution = round(orig_w / (resolution_scale * resolution)), round(orig_h / (resolution_scale * resolution))
    else:  # should be a type that converts to float
        if resolution == -1:
            if orig_w > 1600:
                if not WARNED:
                    print('[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.')
                    print('If this is not desired, please explicitly specify "--resolution/-r" as 1')
                    WARNED = True
                global_down = orig_w / 1600
            else:
                global_down = 1
        else:
            global_down = orig_w / resolution

        scale = float(global_down) * float(resolution_scale)
        resolution = (int(orig_w / scale), int(orig_h / scale))

    resized_image_rgb = PILtoTorch(cam_info.image, resolution)
    gt_image = resized_image_rgb[:3]
    loaded_mask = None
    if resized_image_rgb.shape[1] == 4:
        loaded_mask = resized_image_rgb[-1:]
    
    freq_images = None
    if training:
        if hp.freq_split:
            if hp.split_method == 'linear_split':
                # freq_images = soft_split_numpy(gt_image, hp.d, hp.D)
                freq_images = soft_linear_split_freq(gt_image, hp.d, hp.D)
            elif hp.split_method == 'exp_split':
                freq_images = soft_split_freqs_torch(gt_image, hp.scale)

    return Camera(
        uid=id, colmap_id=cam_info.uid, 
        R=cam_info.R, T=cam_info.T, FoVx=cam_info.FovX, FoVy=cam_info.FovY, 
        gt_image=gt_image, mask=loaded_mask, freq_images=freq_images, image_name=cam_info.image_name, 
    )
